import numpy as np
import pandas as pd
import os

class AntColonyOptimizerAll:
    def __init__(self, distances, n_ants, n_best, n_iterations, decay, alpha=1, beta=1):
        self.distances  = distances
        self.pheromone = np.ones(self.distances.shape) / len(distances)
        self.all_inds = range(len(distances))
        self.n_ants = n_ants
        self.n_best = n_best
        self.n_iterations = n_iterations
        self.decay = decay
        self.alpha = alpha
        self.beta = beta

# run_all
    def run_all(self, start, end):
        shortest_path = None
        best_length = float('inf')
        for _ in range(self.n_iterations):
            paths = self.generate_paths_all(start, end)
            self.spread_pheromone(paths, self.n_best, shortest_path=shortest_path)
            shortest_path = min(paths, key=lambda x: x[1])
            if shortest_path[1] < best_length:
                best_length = shortest_path[1]
                best_path = shortest_path[0]
            self.pheromone *= self.decay
        return best_path, best_length

# spread_pheromone
    def spread_pheromone(self, paths, n_best, shortest_path):
        sorted_paths = sorted(paths, key=lambda x: x[1])
        for path, length in sorted_paths[:n_best]:
            for move in path:
                self.pheromone[move] += 1.0 / self.distances[move]
    
# genrate_paths_all
    def generate_paths_all(self, start, end):
        paths = []
        for _ in range(self.n_ants):
            path = [start]
            while path[-1] != end:
                move = self.possible_move_all(path)
                path.append(move)
            paths.append((path, self.path_length(path)))
        return paths

# possible_move_all
    def possible_move_all(self, path):
        moves = list(set(self.all_inds) - set(path))
        pheromone = np.power(self.pheromone[path[-1]][moves], self.alpha)
        distance = np.reciprocal(np.power(self.distances[path[-1]][moves], self.beta))
        move_prob = pheromone * distance

        total_prob = np.sum(move_prob)
        if total_prob > 0:
            move_prob /= total_prob
        else:
            move_prob = np.ones(len(moves)) / len(moves)  # 平均分配概率以避免 NaN

        next_move = np.random.choice(moves, 1, p=move_prob)[0]
        return next_move

# path_length
    def path_length(self, path):
        length = 0
        for i in range(len(path) - 1):
            length += self.distances[path[i]][path[i+1]]
        return length


##############################################################################

# # 测试项目目录
# cwd = os.getcwd()
# print(f"Current working directory: {cwd}")

# # Define the distances matrix
# distances = np.array([
#     [0, 2, 3],
#     [2, 0, 4],
#     [3, 4, 0]
# ])

# # file_path = 'map_matrix.csv'
# # if os.path.exists(file_path):
# #     print('File exists')
# #     df = pd.read_csv(file_path)
# #     distances = df.values
# # else:
# #     print('File does not exist')

# ##############################################################################

# print("##############################################################################")

# ##############################################################################

# # Parameters
# n_ants = 5
# n_best = 1
# n_iterations = 100
# decay = 0.95
# alpha = 1  # Pheromone importance
# beta = 1   # Distance importance

# ##############################################################################

# # Running the ACO all
# aco_all = AntColonyOptimizerAll(distances, n_ants, n_best, n_iterations, decay, alpha, beta)
# start_node = 0
# end_node = 2
# best_path, best_length = aco_all.run_all(start_node, end_node)
# print("All Best path:", best_path)
# print("All Best path length:", best_length)

# ##############################################################################

# print("##############################################################################")

# ##############################################################################